import os
import sys
from pathlib import Path
from cbs import CBSSolver
import json
from PyQt5.QtWidgets import QApplication, QMainWindow, QWidget, QVBoxLayout, QLabel, QPushButton, QMessageBox
from PyQt5.QtCore import Qt, QPoint
from single_agent_planner import get_sum_of_cost  # 添加这个导入
from visualize import Animation3D  # 添加这个导入


class HoverMoveButton(QPushButton):
    def __init__(self, parent=None):
        super().__init__(parent)
        self.original_pos = self.pos()
        self.move_offset = QPoint(5, 5)  # 向右下角移动的偏移量

    def enterEvent(self, event):
        self.original_pos = self.pos()
        self.move(self.pos() + self.move_offset)
        super().enterEvent(event)

    def leaveEvent(self, event):
        self.move(self.original_pos)
        super().leaveEvent(event)


def print_mapf_instance(my_map, starts, goals):
    """打印地图和位置信息"""
    print('Start locations')
    print_locations(my_map, starts)
    print('Goal locations')
    print_locations(my_map, goals)


def print_locations(my_map, locations):
    """在地图上标记位置"""
    starts_map = [[-1 for _ in range(len(my_map[0]))] for _ in range(len(my_map))]
    for i in range(len(locations)):
        starts_map[locations[i][0]][locations[i][1]] = i
    to_print = ''
    for x in range(len(my_map)):
        for y in range(len(my_map[0])):
            if starts_map[x][y] >= 0:
                to_print += str(starts_map[x][y]) + ' '
            elif my_map[x][y]:
                to_print += '@ '
            else:
                to_print += '. '
        to_print += '\n'
    print(to_print)


def import_mapf_instance(filename):
    """从文件中导入地图和代理的起始及目标位置"""
    try:
        with open(filename, 'r', encoding='utf-8') as f:
            # 第一行读取矩阵大小
            rows, cols = map(int, f.readline().strip().split())

            # 读取矩阵数据
            my_map = []
            for _ in range(rows):
                line = f.readline().strip()
                # 处理每行数据，确保正确解析空格分隔的字符
                row = [c == '@' for c in line.split()]
                my_map.append(row)

            # 读取点对数量
            num_agents = int(f.readline().strip())

            # 读取起终点位置
            starts = []
            goals = []
            for _ in range(num_agents):
                line = f.readline().strip()
                if line:
                    x1, y1, x2, y2 = map(int, line.split())
                    starts.append((x1, y1))
                    goals.append((x2, y2))

            if not starts or not goals:
                raise ValueError("没有找到有效的起点和终点")

            return my_map, starts, goals

    except Exception as e:
        raise Exception(f"处理文件 {filename} 时出错: {str(e)}")


def save_animation_data(original_file, my_map, starts, goals, paths):
    """保存完整动画数据到JSON文件"""
    try:
        # 确保数据格式正确
        data = {
            "map": [[bool(cell) for cell in row] for row in my_map],
            "starts": [list(start) for start in starts],
            "goals": [list(goal) for goal in goals],
            "paths": [[list(pos) for pos in path] for path in paths]
        }

        # 获取象限编号
        quadrant_num = os.path.basename(original_file).split('_')[1].split('.')[0]

        # 保存路径
        current_dir = os.path.dirname(os.path.abspath(__file__))
        data_dir = os.path.join(current_dir, "data")
        os.makedirs(data_dir, exist_ok=True)
        save_path = os.path.join(data_dir, f"quadrant_{quadrant_num}.animation.json")

        # 保存为JSON
        with open(save_path, 'w', encoding='utf-8') as f:
            json.dump(data, f, indent=2)
        print(f"已保存动画数据到: {save_path}")
        return save_path
    except Exception as e:
        print(f"保存动画数据时出错: {str(e)}")
        return None


def load_animation_data(json_file):
    """从JSON文件加载动画数据"""
    with open(json_file, encoding='utf-8') as f:
        data = json.load(f)
    return (
        [[bool(cell) for cell in row] for row in data['map']],  # 转换为布尔数组
        [tuple(s) for s in data['starts']],  # 转换为元组
        [tuple(g) for g in data['goals']],  # 转换为元组
        [[tuple(p) for p in path] for path in data['paths']]  # 转换为元组列表
    )


def run_path_planning(instance_file, solver_type="CBS"):
    """运行路径规划"""
    try:
        my_map = []
        starts = []
        goals = []

        # 读取地图和点对数据
        with open(instance_file, 'r', encoding='utf-8') as f:
            # 读取矩阵大小
            rows, cols = map(int, f.readline().strip().split())

            # 读取矩阵数据
            for _ in range(rows):
                line = f.readline().strip()
                row = [c == '@' for c in line.split()]
                my_map.append(row)

            # 首先保存基本的地图数据
            base_data = {
                "map": [[bool(cell) for cell in row] for row in my_map],
                "starts": [],
                "goals": [],
                "paths": []
            }

            # 获取象限编号并保存基本数据
            quadrant_num = os.path.basename(instance_file).split('_')[1].split('.')[0]
            current_dir = os.path.dirname(os.path.abspath(__file__))
            data_dir = os.path.join(current_dir, "data")
            os.makedirs(data_dir, exist_ok=True)
            animation_file = os.path.join(data_dir, f"quadrant_{quadrant_num}.animation.json")

            with open(animation_file, 'w', encoding='utf-8') as af:
                json.dump(base_data, af, indent=2)
            print(f"已保存基本动画数据到: {animation_file}")

            # 尝试读取点对数据
            num_agents = int(f.readline().strip())
            for _ in range(num_agents):
                line = f.readline().strip()
                if line:
                    x1, y1, x2, y2 = map(int, line.split())
                    starts.append((x1, y1))
                    goals.append((x2, y2))

        # 如果没有点对，直接返回
        if not starts or not goals:
            return {"paths": [], "map": my_map}, 0, 0

        total_pairs = len(starts)

        # 验证点对
        valid_pairs = []
        for i, (start, goal) in enumerate(zip(starts, goals)):
            if my_map[start[0]][start[1]] or my_map[goal[0]][goal[1]]:
                print(f"跳过无效的路径点对 {i+1}")
                continue
            valid_pairs.append((start, goal))

        if not valid_pairs:
            print("没有找到有效的点对")
            return {"paths": [], "map": my_map}, total_pairs, 0

        valid_starts, valid_goals = zip(*valid_pairs)

        # CBS规划
        if solver_type == "CBS":
            solver = CBSSolver(my_map, valid_starts, valid_goals)
            try:
                paths = solver.find_solution()
                if not paths:
                    print(f"无法找到有效路径组合")
                    return {"paths": [], "map": my_map}, total_pairs, 0

                # 更新动画数据，包含路径信息
                full_data = {
                    "map": [[bool(cell) for cell in row] for row in my_map],
                    "starts": [[int(x), int(y)] for x, y in valid_starts],
                    "goals": [[int(x), int(y)] for x, y in valid_goals],  # 修复这行的语法错误
                    "paths": [[[int(x), int(y)] for x, y in path] for path in paths]
                }

                # 保存完整的动画数据
                with open(animation_file, 'w', encoding='utf-8') as f:
                    json.dump(full_data, f, indent=2)
                print(f"已保存完整动画数据到: {animation_file}")

                # 计算成本
                cost = get_sum_of_cost(paths)

                result = {
                    'paths': paths,
                    'starts': valid_starts,
                    'goals': valid_goals,
                    'map': my_map
                }

                return result, total_pairs, cost

            except BaseException as be:
                print(f"路径规划失败: {str(be)}")
                return {"paths": [], "map": my_map}, total_pairs, 0

        else:
            raise ValueError(f"不支持的求解器类型: {solver_type}")

    except Exception as e:
        print(f"处理文件出错: {str(e)}")
        return {"paths": [], "map": my_map}, 0, 0


def show_animation(saved_file):
    """显示3D动画"""
    if not saved_file:
        return

    try:
        my_map, starts, goals, paths = load_animation_data(saved_file)
        animation = Animation3D(my_map, starts, goals, paths)
        animation.show()

    except Exception as e:
        QMessageBox.critical(None, "错误", f"显示动画时出错: {str(e)}")


def read_num_pairs():
    """
    从 point_pairs.txt 文件中读取点对数量
    """
    try:
        current_dir = os.path.dirname(os.path.abspath(__file__))
        file_path = os.path.join(current_dir, "data", "point_pairs.txt")
        with open(file_path, 'r') as file:
            first_line = file.readline().strip()
            return int(first_line)
    except (FileNotFoundError, ValueError):
        QMessageBox.critical(None, "错误", "无法从 point_pairs.txt 文件中读取有效的点对数量。")
        return None


def get_path_cost(path):
    """计算单条路径的成本（长度）"""
    if not path:
        return 0
    cost = 0
    for i in range(1, len(path)):
        # 计算两点间的曼哈顿距离
        x1, y1 = path[i - 1]
        x2, y2 = path[i]
        cost += abs(x2 - x1) + abs(y2 - y1)
    return cost


class MainWindow(QMainWindow):
    def __init__(self):
        super().__init__()
        self.solver = None
        self.initUI()

    def initUI(self):
        self.setWindowTitle('路径规划实验')
        self.setGeometry(100, 100, 800, 600)

        central_widget = QWidget()
        self.setCentralWidget(central_widget)
        layout = QVBoxLayout(central_widget)

        self.status_label = QLabel('准备开始路径规划...', self)
        layout.addWidget(self.status_label)

        self.start_button = HoverMoveButton('开始规划', self)
        self.start_button.clicked.connect(self.run_planning)
        layout.addWidget(self.start_button)

    def run_planning(self):
        """执行路径规划并返回是否成功"""
        try:
            current_dir = os.path.dirname(os.path.abspath(__file__))
            data_dir = os.path.join(current_dir, "data")

            # 读取组合文件
            input_file = os.path.join(data_dir, "combined_output.txt")
            with open(input_file, 'r') as f:
                # 读取层数
                num_layers = int(f.readline())

                # 读取矩阵大小
                rows, cols = map(int, f.readline().split())
                # 读取并转换矩阵
                my_map = []
                for _ in range(rows):
                    line = f.readline().strip()
                    my_map.append([c == '@' for c in line])  # 转换为布尔数组

                # 读取点对数量和点对
                num_pairs = int(f.readline())
                starts = []
                goals = []
                for _ in range(num_pairs):
                    x1, y1, x2, y2 = map(int, f.readline().split())
                    starts.append((x1, y1))
                    goals.append((x2, y2))

            # 创建 CBS 求解器实例
            self.solver = CBSSolver(my_map, starts, goals)

            # 执行路径规划
            paths = self.solver.find_solution()

            if paths:
                # 创建结果数据结构
                result = {
                    'matrix': {
                        'rows': rows,
                        'cols': cols,
                        'data': [[not cell for cell in row] for row in my_map]
                    },
                    'paths': []
                }

                # 添加每条路径的信息
                for i, path in enumerate(paths):
                    path_info = {
                        'start': list(starts[i]),
                        'goal': list(goals[i]),
                        'path': [list(pos) for pos in path]
                    }
                    result['paths'].append(path_info)

                # 保存结果到 JSON 文件
                output_file = os.path.join(data_dir, "planning_results.json")
                with open(output_file, 'w') as f:
                    json.dump(result, f, indent=4)

                return True

            return False

        except Exception as e:
            print(f"路径规划失败: {str(e)}")
            return False

    def show_animation(self):
        """显示3D动画"""
        if not self.solver:
            return

        # 获取最新的规划结果
        current_dir = os.path.dirname(os.path.abspath(__file__))
        results_file = os.path.join(current_dir, "data", "planning_results.json")

        try:
            with open(results_file, 'r') as f:
                results = json.load(f)

            # 创建并显示动画
            animation = Animation3D(
                [results['matrix']['data']],  # 地图列表
                [[tuple(p['start']) for p in results['paths']]],  # 起点列表
                [[tuple(p['goal']) for p in results['paths']]],  # 终点列表
                [[tuple(pos) for pos in p['path']] for p in results['paths']]  # 路径列表
            )
            animation.show()

        except Exception as e:
            print(f"显示动画时出错: {str(e)}")


if __name__ == '__main__':
    app = QApplication(sys.argv)
    window = MainWindow()
    window.show()
    sys.exit(app.exec_())